#!/usr/bin/env python3
import os
import subprocess
import sys
import time
from datetime import datetime

# Configuration
REMOTE = '//192.168.31.57/raid5'
MOUNT_POINT = '/mnt/raid5'
CREDS_FILE = '/root/.smb-cred-raid5'
USERNAME = 'harryshen'
PASSWORD = 'zqbx66666666'

# Color codes for better output
class Colors:
    GREEN = '\033[92m'
    RED = '\033[91m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    PURPLE = '\033[95m'
    CYAN = '\033[96m'
    WHITE = '\033[97m'
    BOLD = '\033[1m'
    END = '\033[0m'

def print_success(message):
    print(f"{Colors.GREEN}{Colors.BOLD}✓ {message}{Colors.END}")

def print_error(message):
    print(f"{Colors.RED}{Colors.BOLD}✗ {message}{Colors.END}")

def print_info(message):
    print(f"{Colors.BLUE}{Colors.BOLD}ℹ {message}{Colors.END}")

def print_warning(message):
    print(f"{Colors.YELLOW}{Colors.BOLD}⚠ {message}{Colors.END}")

def print_step(message):
    print(f"{Colors.CYAN}{Colors.BOLD}→ {message}{Colors.END}")

def run(cmd, check=True, silent=False):
    """Run a command and return the output"""
    if not silent:
        print_step(f"Executing: {cmd}")
    
    result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    returncode = result.returncode
    output = result.stdout.strip()
    
    if check and returncode != 0:
        print_error(f"Command failed: {cmd}")
        print_error(f"Output: {output}")
        sys.exit(returncode)
    
    if not silent and output:
        print(f"  {output}")
    
    return output


def ensure_tools():
    """Ensure required tools are installed"""
    print_info("Checking for required tools...")
    
    have_cifs = run('which mount.cifs || which mount | grep -q cifs && echo ok || echo missing', check=False, silent=True).endswith('ok')
    have_smb = run('which smbclient >/dev/null 2>&1 && echo ok || echo missing', check=False, silent=True).endswith('ok')
    
    if not (have_cifs and have_smb):
        print_warning("Installing required packages...")
        run('apt-get update -y && apt-get install -y cifs-utils smbclient', silent=False)
        print_success("Required packages installed successfully")
    else:
        print_success("All required tools are available")


def ensure_mount_point():
    """Ensure mount point directory exists"""
    print_info(f"Checking mount point: {MOUNT_POINT}")
    
    if not os.path.isdir(MOUNT_POINT):
        print_step(f"Creating mount point directory: {MOUNT_POINT}")
        run(f'mkdir -p {MOUNT_POINT}')
        print_success(f"Mount point created: {MOUNT_POINT}")
    else:
        print_success(f"Mount point already exists: {MOUNT_POINT}")


def write_credentials():
    """Write SMB credentials to secure file"""
    print_info("Setting up credentials file...")
    
    try:
        # Create credentials file directly with Python
        with open(CREDS_FILE, 'w') as f:
            f.write(f"username={USERNAME}\n")
            f.write(f"password={PASSWORD}\n")
        
        # Set proper permissions
        os.chmod(CREDS_FILE, 0o600)
        
        print_success(f"Credentials file created: {CREDS_FILE}")
        
        # Verify the file was created correctly
        with open(CREDS_FILE, 'r') as f:
            content = f.read()
            if USERNAME in content and PASSWORD in content:
                print_success("Credentials file verified successfully")
            else:
                print_error("Credentials file verification failed")
                
    except Exception as e:
        print_error(f"Failed to create credentials file: {str(e)}")
        sys.exit(1)


def is_already_mounted():
    """Check if the share is already mounted"""
    print_info("Checking if share is already mounted...")
    
    mounts = run('mount | grep -w cifs | cat', check=False, silent=True)
    is_mounted = REMOTE in mounts or MOUNT_POINT in mounts
    
    if is_mounted:
        print_warning("Share is already mounted")
    else:
        print_info("Share is not currently mounted")
    
    return is_mounted


def test_connection():
    """Test SMB connection before mounting"""
    print_info("Testing SMB connection...")
    
    try:
        # Test with smbclient
        test_cmd = f'smbclient -L {REMOTE} -U {USERNAME}%{PASSWORD} -N 2>/dev/null | head -5'
        result = run(test_cmd, check=False, silent=True)
        
        if result and "Sharename" in result:
            print_success("SMB connection test successful")
            return True
        else:
            print_warning("SMB connection test inconclusive, proceeding with mount attempt")
            return True
    except:
        print_warning("SMB connection test failed, proceeding with mount attempt")
        return True


def mount_share():
    """Mount the SMB share"""
    print_info(f"Mounting {REMOTE} to {MOUNT_POINT}")
    
    opts = f'credentials={CREDS_FILE},iocharset=utf8,vers=3.0,uid=0,gid=0,file_mode=0777,dir_mode=0777'
    run(f'mount -t cifs {REMOTE} {MOUNT_POINT} -o {opts}')
    
    # Verify mount was successful
    time.sleep(1)
    if os.path.ismount(MOUNT_POINT):
        print_success(f"Successfully mounted {REMOTE} at {MOUNT_POINT}")
        return True
    else:
        print_error(f"Failed to mount {REMOTE} at {MOUNT_POINT}")
        return False


def show_mount_info():
    """Display mount information and directory contents"""
    print_info("Mount Information:")
    
    # Show mount details
    mount_info = run(f'mount | grep {MOUNT_POINT}', check=False, silent=True)
    if mount_info:
        print(f"  {mount_info}")
    
    # Show disk usage
    try:
        df_info = run(f'df -h {MOUNT_POINT}', check=False, silent=True)
        if df_info:
            print(f"  Disk usage: {df_info.split('\\n')[-1]}")
    except:
        pass
    
    # Show directory contents
    print_info("Directory contents:")
    try:
        contents = run(f'ls -la {MOUNT_POINT} | head -20', check=False, silent=True)
        if contents:
            for line in contents.split('\\n'):
                print(f"  {line}")
    except:
        print_warning("Could not list directory contents")


def main():
    """Main function to mount RAID5 SMB share"""
    print(f"{Colors.PURPLE}{Colors.BOLD}{'='*60}{Colors.END}")
    print(f"{Colors.PURPLE}{Colors.BOLD}  RAID5 SMB Mount Script{Colors.END}")
    print(f"{Colors.PURPLE}{Colors.BOLD}  Started at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}{Colors.END}")
    print(f"{Colors.PURPLE}{Colors.BOLD}{'='*60}{Colors.END}")
    
    # Check if running as root
    if os.geteuid() != 0:
        print_error('This script must be run with sudo or as root.')
        print_info('Please run: sudo python3 mount_raid5.py')
        sys.exit(1)
    
    try:
        # Step 1: Ensure required tools are installed
        ensure_tools()
        
        # Step 2: Ensure mount point exists
        ensure_mount_point()
        
        # Step 3: Write credentials file
        write_credentials()
        
        # Step 4: Test connection
        test_connection()
        
        # Step 5: Check if already mounted
        if is_already_mounted():
            print_success(f'Share is already mounted: {REMOTE} at {MOUNT_POINT}')
        else:
            # Step 6: Mount the share
            if mount_share():
                print_success(f'Successfully mounted {REMOTE} at {MOUNT_POINT}')
            else:
                print_error('Failed to mount the share')
                sys.exit(1)
        
        # Step 7: Show mount information
        show_mount_info()
        
        # Final success message
        print(f"{Colors.GREEN}{Colors.BOLD}{'='*60}{Colors.END}")
        print(f"{Colors.GREEN}{Colors.BOLD}  ✓ MOUNT OPERATION COMPLETED SUCCESSFULLY! ✓{Colors.END}")
        print(f"{Colors.GREEN}{Colors.BOLD}  Remote: {REMOTE}{Colors.END}")
        print(f"{Colors.GREEN}{Colors.BOLD}  Mount Point: {MOUNT_POINT}{Colors.END}")
        print(f"{Colors.GREEN}{Colors.BOLD}  Completed at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}{Colors.END}")
        print(f"{Colors.GREEN}{Colors.BOLD}{'='*60}{Colors.END}")
        
    except KeyboardInterrupt:
        print_error("Operation cancelled by user")
        sys.exit(1)
    except Exception as e:
        print_error(f"Unexpected error: {str(e)}")
        sys.exit(1)

if __name__ == '__main__':
    main()
